#include <cuda_runtime.h>
#include <opencv2/opencv.hpp>
#include <stdio.h>
#include <algorithm>
#include <chrono>

using namespace cv;

#define min(a, b)   ((a) < (b) ? (a) : (b))
#define checkRuntime(op) __check_cuda_runtime((op), #op, __FILE__, __LINE__)

double timestamp_now_float() {
    return std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::system_clock::now().time_since_epoch()).count() / 1000.0;
}

bool __check_cuda_runtime(cudaError_t code, const char* op, const char* file, int line){
    if(code != cudaSuccess){
        const char* err_name = cudaGetErrorName(code);
        const char* err_message = cudaGetErrorString(code);
        printf("runtime error %s:%d %s failed. \n code = %s, message = %s\n", file, line, op, err_name, err_message);
        return false;
    }
    return true;
}

void warp_affine_bilinear(// 声明
    uint8_t* src, int src_line_size, int src_width, int src_height,
    uint8_t* dst, int dst_line_size, int dst_width, int dst_height,
    uint8_t fill_value
);

Mat warpaffine_to_center_align(Mat image, const Size& size){
    // 步骤1 在CPU上开辟空间，用于保存预处理完成后的图像
    Mat output(size, CV_8UC3);

    // 步骤2 在GPU上开辟空间，将读取的图像数据拷贝到GPU上
    uint8_t* psrc_device = nullptr;
    uint8_t* pdst_device = nullptr;

    size_t src_size = image.cols * image.rows * 3;
    size_t dst_size = size.width * size.height * 3;

    checkRuntime(cudaMalloc(&psrc_device, src_size));
    checkRuntime(cudaMalloc(&pdst_device, dst_size));

    checkRuntime(cudaMemcpy(psrc_device, image.data, src_size, cudaMemcpyHostToDevice));
    
    // 步骤3 调用预处理函数
    warp_affine_bilinear(// 调用
        psrc_device, image.cols * 3, image.cols, image.rows,
        pdst_device, size.width * 3, size.width, size.height,
        114
    );

    // 步骤4 将预处理后的图片重新拷贝到CPU上进行保存
    checkRuntime(cudaPeekAtLastError());    // 检查核函数执行是否存在错误
    checkRuntime(cudaMemcpy(output.data, pdst_device, dst_size, cudaMemcpyDeviceToHost));

    // 步骤5 释放CPU和GPU上分配的内存
    checkRuntime(cudaFree(pdst_device));
    checkRuntime(cudaFree(psrc_device));

    return output;
}

int main(){
    auto t0 = timestamp_now_float();
    Mat image = imread("cat.png");
    Mat output = warpaffine_to_center_align(image, Size(640, 640));
    imwrite("output.png", output);
    auto fee = timestamp_now_float() - t0;
    printf("Done. save to output.png fee = %.2f ms\n", fee);
    
    return 0;
}